import os
import cv2
import scipy.io as scio
import h5py

def dataloader(source_dir, gt_dir, dataset):
    file_list = sorted(os.listdir(source_dir))
    res = []

    for file_name in file_list:
        name = file_name.split('.')[0]
        gt_name = name + '.mat'

        image_path = os.path.join(source_dir, file_name)
        gt_path = os.path.join(gt_dir, gt_name)

        image = cv2.imread(image_path, 1)

        if dataset == 'whsymmax':
            gt = h5py.File(gt_path)
        else:
            gt = scio.loadmat(gt_path)

        if dataset== 'sklarge' or dataset == 'sympascal' or dataset == 'sk506':
            gt = gt['symmetry']
        elif dataset == 'whsymmax':
            gt = gt['sym'].value
            gt = gt.transpose((1, 0))
        elif dataset == 'symmax300':
            gt = gt['gt']
        gt[gt > 0] = 1

        res.append((image, gt, name))
    
    return res


def data_filter(data_list, path):
    res = []
    for index in range(len(data_list)):
        _, _, name = data_list[index]
        target_path = os.path.join(path, name + 'png')
        if not os.path.exists(target_path):
            res.append(data_list[index])
    return res

